# 实体识别agent

# entity_recognizer.py
import re
import jieba
import jieba.posseg as pseg
from typing import Dict, List, Any, Optional, Tuple
from datetime import datetime, timedelta
import logging

class EntityRecognizer:
    """查询实体识别器 - 支持完整NL2SQL流程"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self._initialize_chinese_dict()
        
        # 支持多数据库的表列模式
        self.database_patterns = {
            "中国城市": {
                "tables": ["城市", "省份", "省生产总值", "市生产总值", "中国城市称号"],
                "columns": ["城市id", "名称", "面积", "人口", "AQ", "市花", "所属省id", 
                           "省份id", "省份名称", "省面积", "省人口", "生产总值", "GDP", 
                           "年度", "经济指标", "气象条件", "备注", "称号", "荣誉"]
            },
            "X公司工资": {
                "tables": ["工资表", "员工表", "薪资表", "部门表"],
                "columns": ["工资", "基本工资", "奖金", "补贴", "实发工资", "员工编号", "部门", "发放日期"]
            }
        }
        
        # 查询意图关键词
        self.intent_keywords = {
            "instructional": ["显示", "查询", "查找", "检索", "获取", "查看", "列出"],
            "exploratory": ["分析", "统计", "趋势", "分布", "对比", "排名", "TOP", "前"],
            "data_extraction": ["导出", "下载", "报表", "报告", "全部", "所有", "详细"]
        }
        
        # 时间表达式模式
        self.time_patterns = [
            (r'(\d+)年前', 'years_ago'),
            (r'(\d+)个月前', 'months_ago'), 
            (r'(\d+)天前', 'days_ago'),
            (r'(\d+)周前', 'weeks_ago'),
            (r'去年', 'last_year'),
            (r'今年', 'this_year'),
            (r'明年', 'next_year'),
            (r'上个月', 'last_month'),
            (r'本月', 'this_month'),
            (r'下个月', 'next_month'),
            (r'上周', 'last_week'),
            (r'本周', 'this_week'),
            (r'下周', 'next_week'),
            (r'昨天', 'yesterday'),
            (r'今天', 'today'),
            (r'明天', 'tomorrow'),
            (r'最近', 'recent'),
            (r'近期', 'recent'),
            (r'第一季度', 'q1'),
            (r'第二季度', 'q2'),
            (r'第三季度', 'q3'),
            (r'第四季度', 'q4')
        ]
    
    def _initialize_chinese_dict(self):
        """初始化中文分词词典"""
        terms = [
            # 数据库表相关
            '中国城市', '城市表', '省份表', '省生产总值', '市生产总值', '中国城市称号',
            '工资表', '员工表', '薪资表', '部门表', '销售表', '订单表', '产品表',
            # 列名相关
            '城市id', '省份id', '所属省id', '省面积', '省人口', '生产总值', 'GDP',
            '基本工资', '实发工资', '员工编号', '发放日期', '销售金额', '订单数量',
            # 业务术语
            '北京市', '天津市', '上海市', '重庆市', '河北省', '广东省', '浙江省'
        ]
        
        for term in terms:
            jieba.add_word(term, freq=1000, tag='n')
    
    def recognize_entities(self, query: str, database_name: str = None) -> Dict[str, Any]:
        """
        从自然语言查询中识别实体
        
        Args:
            query: 自然语言查询
            database_name: 目标数据库名称
            
        Returns:
            实体识别结果，包含后续流程所需的所有信息
        """
        self.logger.info(f"开始实体识别: {query}")
        
        entities = {
            "original_query": query,
            "database": database_name,
            # 核心实体
            "tables": [],
            "columns": [],
            "values": [],
            "conditions": [],
            "time_conditions": [],
            "aggregations": [],
            "relationships": [],
            # 意图识别相关
            "intent_clues": {
                "instructional_keywords": [],
                "exploratory_keywords": [], 
                "extraction_keywords": []
            },
            # 查询复杂度评估
            "complexity": "simple",
            # 置信度
            "confidence": "medium",
            "processing_time": None
        }
        
        start_time = time.time()
        
        try:
            # 分词处理
            words = pseg.cut(query)
            word_list = list(words)
            
            # 识别查询意图线索
            entities["intent_clues"] = self._recognize_intent_clues(query)
            
            # 识别表名
            tables = self._recognize_tables(query, word_list, database_name)
            entities["tables"] = tables
            
            # 识别列名  
            columns = self._recognize_columns(query, word_list, database_name, tables)
            entities["columns"] = columns
            
            # 识别值条件
            values = self._recognize_values(query, word_list)
            entities["values"] = values
            
            # 识别时间条件
            time_conditions = self._recognize_time_conditions(query)
            entities["time_conditions"] = time_conditions
            
            # 识别聚合操作
            aggregations = self._recognize_aggregations(query)
            entities["aggregations"] = aggregations
            
            # 识别条件关系
            conditions = self._recognize_conditions(query)
            entities["conditions"] = conditions
            
            # 评估查询复杂度
            entities["complexity"] = self._assess_complexity(entities)
            
            # 计算置信度
            entities["confidence"] = self._calculate_confidence(entities)
            
            entities["processing_time"] = time.time() - start_time
            
            self.logger.info(f"实体识别完成: {len(tables)}表, {len(columns)}列, 复杂度: {entities['complexity']}")
            return entities
            
        except Exception as e:
            self.logger.error(f"实体识别失败: {e}")
            entities["confidence"] = "low"
            entities["error"] = str(e)
            entities["processing_time"] = time.time() - start_time
            return entities
    
    def _recognize_intent_clues(self, query: str) -> Dict[str, List]:
        """识别查询意图线索"""
        intent_clues = {
            "instructional_keywords": [],
            "exploratory_keywords": [],
            "extraction_keywords": []
        }
        
        query_lower = query.lower()
        
        # 识别指令式查询关键词
        for keyword in self.intent_keywords["instructional"]:
            if keyword in query_lower:
                intent_clues["instructional_keywords"].append(keyword)
        
        # 识别探索式分析关键词
        for keyword in self.intent_keywords["exploratory"]:
            if keyword in query_lower:
                intent_clues["exploratory_keywords"].append(keyword)
        
        # 识别数据提取关键词
        for keyword in self.intent_keywords["data_extraction"]:
            if keyword in query_lower:
                intent_clues["extraction_keywords"].append(keyword)
        
        return intent_clues
    
    def _recognize_tables(self, query: str, words: List, database_name: str) -> List[Dict]:
        """识别表名"""
        tables = []
        
        # 基于数据库特定模式匹配
        if database_name in self.database_patterns:
            db_patterns = self.database_patterns[database_name]
            for table_name in db_patterns["tables"]:
                if self._is_entity_mentioned(table_name, query):
                    table_info = {
                        "name": table_name,
                        "confidence": "high",
                        "source": "pattern_match",
                        "matched_text": table_name
                    }
                    tables.append(table_info)
        
        # 从分词结果中识别
        for word, flag in words:
            if flag in ['n', 'vn'] and len(word) >= 2:
                if self._is_likely_table_name(word) and not any(t['name'] == word for t in tables):
                    table_info = {
                        "name": word,
                        "confidence": "low", 
                        "source": "word_segmentation",
                        "matched_text": word
                    }
                    tables.append(table_info)
        
        return tables
    
    def _recognize_columns(self, query: str, words: List, database_name: str, tables: List[Dict]) -> List[Dict]:
        """识别列名"""
        columns = []
        
        # 基于数据库特定模式匹配
        if database_name in self.database_patterns:
            db_patterns = self.database_patterns[database_name]
            for column_name in db_patterns["columns"]:
                if self._is_entity_mentioned(column_name, query):
                    # 推断所属表
                    table_name = self._infer_table_for_column(column_name, tables, database_name)
                    column_info = {
                        "name": column_name,
                        "table": table_name,
                        "confidence": "high",
                        "source": "pattern_match",
                        "matched_text": column_name
                    }
                    columns.append(column_info)
        
        return columns
    
    def _recognize_values(self, query: str, words: List) -> List[Dict]:
        """识别值条件"""
        values = []
        
        # 识别数字值
        number_patterns = [
            r'(\d+\.?\d*)[万亿]?',
            r'(\d+)[万亿]?'
        ]
        
        for pattern in number_patterns:
            numbers = re.findall(pattern, query)
            for num in numbers:
                if float(num.replace('万', '0000').replace('亿', '00000000')) > 1:
                    value_info = {
                        "value": num,
                        "type": "number",
                        "operator": self._infer_operator_for_value(num, query),
                        "confidence": "high"
                    }
                    values.append(value_info)
        
        # 识别字符串值
        quoted_strings = re.findall(r'["\'](.*?)["\']', query)
        for string in quoted_strings:
            if len(string) > 1:
                value_info = {
                    "value": string,
                    "type": "string", 
                    "operator": "=",
                    "confidence": "high"
                }
                values.append(value_info)
        
        return values
    
    def _recognize_time_conditions(self, query: str) -> List[Dict]:
        """识别时间条件"""
        time_conditions = []
        
        for pattern, condition_type in self.time_patterns:
            matches = re.finditer(pattern, query)
            for match in matches:
                time_condition = {
                    "type": condition_type,
                    "matched_text": match.group(0),
                    "value": self._calculate_time_value(condition_type, match),
                    "confidence": "high"
                }
                time_conditions.append(time_condition)
        
        return time_conditions
    
    def _recognize_aggregations(self, query: str) -> List[Dict]:
        """识别聚合操作"""
        aggregations = []
        
        agg_patterns = {
            r'总和|总计|合计|总额': 'SUM',
            r'平均|均值': 'AVG', 
            r'数量|个数|计数': 'COUNT',
            r'最大|最高': 'MAX',
            r'最小|最低': 'MIN',
            r'排名|排序': 'RANK'
        }
        
        for pattern, agg_type in agg_patterns.items():
            if re.search(pattern, query):
                aggregation = {
                    "type": agg_type,
                    "matched_text": re.search(pattern, query).group(0),
                    "confidence": "medium"
                }
                aggregations.append(aggregation)
        
        return aggregations
    
    def _recognize_conditions(self, query: str) -> List[Dict]:
        """识别条件关系"""
        conditions = []
        
        condition_patterns = [
            (r'大于|超过|高于', '>'),
            (r'小于|低于|不足', '<'),
            (r'等于|是|为', '='),
            (r'大于等于|不小于', '>='),
            (r'小于等于|不大于', '<='),
            (r'包含|包括', 'LIKE'),
            (r'不包含|不包括', 'NOT LIKE'),
            (r'在.*之间', 'BETWEEN')
        ]
        
        for pattern, operator in condition_patterns:
            if re.search(pattern, query):
                condition = {
                    "operator": operator,
                    "matched_text": re.search(pattern, query).group(0),
                    "confidence": "medium"
                }
                conditions.append(condition)
        
        return conditions
    
    def _assess_complexity(self, entities: Dict) -> str:
        """评估查询复杂度"""
        tables_count = len(entities["tables"])
        conditions_count = len(entities["conditions"]) + len(entities["time_conditions"])
        aggregations_count = len(entities["aggregations"])
        
        complexity_score = tables_count + conditions_count + aggregations_count
        
        if complexity_score <= 2:
            return "simple"
        elif complexity_score <= 4:
            return "medium"
        else:
            return "complex"
    
    def _is_entity_mentioned(self, entity_name: str, query: str) -> bool:
        """检查实体是否在查询中被提及"""
        return entity_name in query
    
    def _is_likely_table_name(self, word: str) -> bool:
        """判断词语是否可能是表名"""
        if len(word) < 2:
            return False
        
        table_indicators = ['表', '信息', '数据', '记录', '详情', '列表']
        return any(indicator in word for indicator in table_indicators)
    
    def _infer_table_for_column(self, column_name: str, tables: List[Dict], database_name: str) -> str:
        """为列推断所属表"""
        if not tables:
            return ""
        
        # 基于列名模式推断
        if any(keyword in column_name for keyword in ['城市', '市']):
            for table in tables:
                if '城市' in table["name"]:
                    return table["name"]
        
        if any(keyword in column_name for keyword in ['省份', '省']):
            for table in tables:
                if '省份' in table["name"]:
                    return table["name"]
        
        if any(keyword in column_name for keyword in ['工资', '薪资']):
            for table in tables:
                if any(wage_keyword in table["name"] for wage_keyword in ['工资', '薪资']):
                    return table["name"]
        
        return tables[0]["name"] if tables else ""
    
    def _infer_operator_for_value(self, value: str, query: str) -> str:
        """为值推断操作符"""
        query_lower = query.lower()
        
        if '大于' in query_lower or '超过' in query_lower:
            return ">"
        elif '小于' in query_lower or '低于' in query_lower:
            return "<"
        else:
            return "="
    
    def _calculate_time_value(self, condition_type: str, match) -> Any:
        """计算时间条件的值"""
        now = datetime.now()
        
        if condition_type == 'last_year':
            return (now - timedelta(days=365)).strftime('%Y-%m-%d')
        elif condition_type == 'this_year':
            return now.strftime('%Y')
        elif condition_type == 'last_month':
            return (now - timedelta(days=30)).strftime('%Y-%m-%d')
        elif condition_type == 'this_month':
            return now.strftime('%Y-%m')
        elif condition_type == 'yesterday':
            return (now - timedelta(days=1)).strftime('%Y-%m-%d')
        elif condition_type == 'today':
            return now.strftime('%Y-%m-%d')
        elif condition_type == 'years_ago':
            years = int(match.group(1))
            return (now - timedelta(days=365*years)).strftime('%Y-%m-%d')
        elif condition_type == 'months_ago':
            months = int(match.group(1))
            return (now - timedelta(days=30*months)).strftime('%Y-%m-%d')
        elif condition_type == 'days_ago':
            days = int(match.group(1))
            return (now - timedelta(days=days)).strftime('%Y-%m-%d')
        elif condition_type == 'recent':
            return (now - timedelta(days=7)).strftime('%Y-%m-%d')
        else:
            return None
    
    def _calculate_confidence(self, entities: Dict) -> str:
        """计算实体识别的总体置信度"""
        if not entities["tables"] and not entities["columns"]:
            return "low"
        
        high_confidence_entities = sum(1 for table in entities["tables"] if table["confidence"] == "high")
        high_confidence_entities += sum(1 for column in entities["columns"] if column["confidence"] == "high")
        
        if high_confidence_entities >= 2:
            return "high"
        elif high_confidence_entities >= 1:
            return "medium"
        else:
            return "low"

# 添加time模块导入
import time